import os

from PIL import Image, ImageDraw

from cnn_rotate.utils import file_util


def center_crop(img, crop_size):
    try:
        # 打开图片
        # img = Image.open(image_path)
        width, height = img.size

        # 计算裁剪区域
        left = (width - crop_size) // 2
        top = (height - crop_size) // 2
        right = (width + crop_size) // 2
        bottom = (height + crop_size) // 2

        # 确保裁剪区域有效
        if left < 0 or top < 0 or right > width or bottom > height:
            print(f"图片尺寸 ({width}x{height}) 过小，无法裁剪出 {crop_size}x{crop_size} 的区域")
            return None

        # 裁剪图片
        cropped_img = img.crop((left, top, right, bottom))

        # 转换图片模式（如果需要）
        if cropped_img.mode == 'RGBA':
            cropped_img = cropped_img.convert('RGB')

        return cropped_img
    except Exception as e:
        print(f"处理图片时出错: {e}")
        return None


def circle_to_white(img):
    width, height = img.size
    center_x, center_y = width // 2, height // 2

    min_length = min(width, height)

    # if min_length >= 350:
    #     # 百度验证码 满足
    #     radius = 175
    # else:
    #     # 不满足百度验证码，取最窄
    #     radius = min_length // 2

    radius = min_length // 2
    # 创建与原图大小相同的白色背景
    white_bg = Image.new('RGBA', (width, height), (255, 255, 255, 255))

    # 创建圆形蒙版
    mask = Image.new('L', (width, height), 0)
    draw = ImageDraw.Draw(mask)
    draw.ellipse((center_x - radius, center_y - radius, center_x + radius, center_y + radius), fill=255)

    # 将原图中圆形区域粘贴到白色背景上
    white_bg.paste(img, (0, 0), mask)

    # 保存处理后的图片
    # white_bg.convert("RGB").save(output_path)
    white_bg.convert("RGB")
    return white_bg, radius


def circle_and_crop(img):
    img_temp, radius = circle_to_white(img)
    res_img = center_crop(img_temp, radius*2)
    return res_img


def main():

    input_dir = '../../dataset/baidu/input'
    output_path = '../../dataset/baidu/output'

    jpgs_input = file_util.list_jpg_files(input_dir)

    if not os.path.exists(output_path):
        os.mkdir(output_path)

    for jpg_path in jpgs_input:
        img = Image.open(jpg_path).convert("RGBA")
        res_img = circle_and_crop(img)

        if res_img is None:
            continue

        jpg_name = file_util.get_filename(jpg_path, include_ext=False)
        output_name = f"{output_path}/{jpg_name}_out.jpg"
        res_img.save(output_name)


if __name__ == '__main__':
    main()

